# !/usr/bin/env python
# -*- coding: utf-8 -*-
from node_util import node_arr
import random


# 平衡二叉树
# 定义: 1.每个节点的子节点不能超过2
# 2.每个节点的左右子树高度不能超过1
# 3.每个节点的属性值小于右孩子节点的属性值，大于左孩子节点的属性值
class AVLNode:
    def __init__(self, key):
        self.key = key
        self.height = 1
        self.left = None
        self.right = None


def get_height(node):
    if node is None:
        return 0
    return node.height


def get_balance(node):
    if node is None:
        return 0
    return get_height(node.left) - get_height(node.right)


def find_min(node):
    l = node.left
    r = node.right
    if l is not None:
        if l.height == 1:
            return l
        if l.height > 1:
            return find_min(l)
    if r is not None:
        if r.height == 1:
            return r
        if r.height > 1:
            return find_min(r)
    return node


def rotate_left(z):
    y = z.right
    T2 = y.left

    # 执行左旋
    y.left = z
    z.right = T2

    # 更新节点的高度
    z.height = 1 + max(get_height(z.left), get_height(z.right))
    y.height = 1 + max(get_height(y.left), get_height(y.right))

    return y


def rotate_right(y):
    x = y.left
    T2 = x.right

    # 执行右旋
    x.right = y
    y.left = T2

    # 更新节点的高度
    y.height = 1 + max(get_height(y.left), get_height(y.right))
    x.height = 1 + max(get_height(x.left), get_height(x.right))

    return x


def insert(root, key):
    if root is None:
        return AVLNode(key)

    if key < root.key:
        root.left = insert(root.left, key)
    elif key > root.key:
        root.right = insert(root.right, key)

    # 更新节点的高度
    root.height = 1 + max(get_height(root.left), get_height(root.right))

    # 获取平衡因子
    balance = get_balance(root)

    # 进行旋转操作来恢复平衡
    # 左旋
    if balance > 1 and key < root.left.key:
        return rotate_right(root)
    # 右旋
    if balance < -1 and key > root.right.key:
        return rotate_left(root)
    # 左右双旋
    if balance > 1 and key > root.left.key:
        root.left = rotate_left(root.left)
        return rotate_right(root)
    # 右左双旋
    if balance < -1 and key < root.right.key:
        root.right = rotate_right(root.right)
        return rotate_left(root)

    return root


def delete(root, key):
    if root is None:
        return root

    if key < root.key:
        root.left = delete(root.left, key)
    elif key > root.key:
        root.right = delete(root.right, key)
    else:
        # 节点有一个或没有子节点
        if root.left is None:
            return root.right
        elif root.right is None:
            return root.left

        # 节点有两个子节点，找到右子树的最小节点
        root.key = find_min(root.right).key
        # 删除右子树的最小节点
        root.right = delete(root.right, root.key)

    # 更新节点的高度
    root.height = 1 + max(get_height(root.left), get_height(root.right))

    # 获取平衡因子
    balance = get_balance(root)

    # 进行旋转操作来恢复平衡
    # 左旋
    if balance > 1 and get_balance(root.left) >= 0:
        return rotate_right(root)
    # 右旋
    if balance < -1 and get_balance(root.right) <= 0:
        return rotate_left(root)
    # 左右双旋
    if balance > 1 and get_balance(root.left) < 0:
        root.left = rotate_left(root.left)
        return rotate_right(root)
    # 右左双旋
    if balance < -1 and get_balance(root.right) > 0:
        root.right = rotate_right(root.right)
        return rotate_left(root)

    return root


# 中序遍历查看结果
def inorder_traversal_avl(root):
    if root is not None:
        inorder_traversal_avl(root.left)
        print(f"({root.key}, {get_balance(root)})", end=" ")
        inorder_traversal_avl(root.right)


def print_tree(root):
    height = root.height
    # 1 + 2 + 4 + 8 + ... = 2^n - 1
    # 倒数第一层最大节点数
    count = pow(2, height) - 1
    node_arr = [root]
    for i in range(0, count // 2):
        print("", end="\t")
    print(root.key, end="\n")
    while True:
        tmp_arr = []
        tmp_count = 0
        for node in node_arr:
            if node is None:
                tmp_arr.append(None)
                tmp_arr.append(None)
            else:
                tmp_arr.append(node.left)
                tmp_arr.append(node.right)
                if node.left is not None:
                    tmp_count += 1
                if node.right is not None:
                    tmp_count += 1
        if tmp_count == 0:
            break
        height -= 1
        count = pow(2, height) - 1
        for i in range(0, count // 2):
            print("", end="\t")
        node_arr.clear()
        node_arr = tmp_arr
        for node in node_arr:
            if node is None:
                print("", end="\t")
            else:
                print(node.key, end="\t")
            for i in range(0, count):
                print("", end="\t")
        print("", end="\n")


if __name__ == '__main__':
    # 创建空树
    avl_root = None

    # 插入操作
    keys_to_insert = [50, 30, 70, 20, 40, 60, 80]
    for key in keys_to_insert:
        avl_root = insert(avl_root, key)


    print("中序遍历结果:", end="\n")
    inorder_traversal_avl(avl_root)
    print("", end="\n")
    print_tree(avl_root)

    # 删除操作
    delete_key = 30
    avl_root = delete(avl_root, delete_key)

    print("删除节点 30 后中序遍历结果:", end="\n")
    inorder_traversal_avl(avl_root)
    print("", end="\n")
    print_tree(avl_root)
